from SemanticNet import SemanticNet
import torch
import torch.nn.functional as F
import numpy as np
import os, argparse
import argparse
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from datetime import datetime
from utils.dataSetLoader import test_dataset
import torch.nn.functional as F
import numpy as np
import imageio



torch.backends.cudnn.benchmark = True

pth_path = "./checkpoints/SemanticNet/SemanticNet-99.pth"
testsize = 448
for data_name in ['CAMO','CHAMELEON','COD10K' ,'NC4K']:
    data_path = './Dataset/TestDataset/{}/'.format(data_name)
    save_path = './res/{}/preds/'.format(data_name)
    edge_path = './res/{}/details/'.format(data_name)
    GT_path = './res/{}/masks/'.format(data_name)
    
    model = SemanticNet()
    model.load_state_dict(torch.load(pth_path))
    model.cuda()
    model.eval()

    os.makedirs(save_path, exist_ok=True)
    os.makedirs(edge_path, exist_ok=True)
    os.makedirs(GT_path, exist_ok=True)

    image_root = '{}/Imgs/'.format(data_path)
    gt_root = '{}/GT/'.format(data_path)
    
    test_loader = test_dataset(image_root, gt_root, testsize)

    print("Start Testing {}...".format(data_name))

    for i in range(test_loader.size):
        image, gt, name = test_loader.load_data()
        gt = np.asarray(gt, np.float32)
        gt /= (gt.max() + 1e-8)  
        image = image.cuda()

        classifier, detail_map, res = model(image)
        
        detail_map = F.interpolate(detail_map, size=gt.shape, mode='bilinear', align_corners=False)
        res = F.interpolate(res, size=gt.shape, mode='bilinear', align_corners=False) 
        res = res.sigmoid().data.cpu().numpy().squeeze()
#        res = res.data.cpu().numpy().squeeze()
        res = (res - res.min()) / (res.max() - res.min() + 1e-8)
        imageio.imwrite(save_path+name, (res*255).astype(np.uint8))
        
        detail_map = detail_map.data.cpu().numpy().squeeze()
        detail_map = (detail_map- detail_map.min())/(detail_map.max()-detail_map.min()+1e-8)
        imageio.imwrite(edge_path + name, (detail_map*255).astype(np.uint8))


        imageio.imwrite(GT_path + name, (gt*255).astype(np.uint8))
        if (i+1) % 100==0 or i+1 ==test_loader.size:
            print("{} images has been finished ".format(i))


print("{} Test Finished...".format(data_name))
